import erniebot
from rich import print
import json

# List supported models
# models = erniebot.Model.list()

# print(models)
# ernie-3.5               文心大模型（ernie-3.5）
# ernie-3.5-8k            文心大模型（ernie-3.5-8k）
# ernie-lite              文心大模型（ernie-lite）
# ernie-4.0               文心大模型（ernie-4.0）
# ernie-speed             文心大模型（ernie-speed）
# ernie-speed-128k        文心大模型（ernie-speed-128k）
# ernie-tiny-8k           文心大模型（ernie-tiny-8k）
# ernie-char-8k           文心大模型（ernie-char-8k）
# ernie-text-embedding    文心百中语义模型
# ernie-vilg-v2           文心一格模型

# Set authentication params
# erniebot.api_type = "aistudio"
# erniebot.access_token = "8c1f786a2a69d12d76088175a0d401c160a41cf0"

# # Create a chat completion
# response = erniebot.ChatCompletion.create(model="ernie-3.5", messages=[{"role": "user", "content": "你好，请介绍下你自己"}])

# print(response.get_result())


# messages详情
# `messages`为一个Python list，其中每个元素为一个dict。在如下示例中，为了与模型进行多轮对话，我们将模型的回复结果插入到`messages`中再继续请求：

# [
#     {
#         "role": "user",
#         "content": "你好啊"
#     },
#     {
#         "role": "assistant",
#         "content": "你好，我是文心一言"
#     },
#     {
#         "role": "user",
#         "content": "深圳周末去哪里玩好?"
#     }
# ]
# `messages`中的每个元素包含如下键值对： | 键名 | 值类型 | 必填 | 值描述 | |:--- | :--- | :--- | :--- | | role | str | 是 | `"user"`表示用户，`"assistant"`表示对话助手，`"function"`表示函数。 | | content | str or None | 是 | 当`role`不为`"function"`时，表示消息内容，必须设置此参数为非`None`值；当`role`为`"function"`时，表示函数响应参数，可以设置此参数为`None`。 | | name | str | 否 | 消息的作者。当`role`为`"function"`时，此参数必填，且是`function_call`中的`name`。 | | function_call | dict | 否 | 由模型生成的函数调用信息，包含函数名称和请求参数等。 | `function_call`为一个Python dict，其中包含如下键值对： | 键名 | 值类型 | 必填 | 值描述 | |:--- | :--- | :--- | :--- | | name | str | 是 | 函数名称。 | | thoughts | str | 否 | 模型思考过程。 | | arguments | str | 是 | 请求参数。 |


class ERNIE_API:
    def __init__(self, aistudio_access_token, system_prompt, function_map):
        erniebot.api_type = "aistudio"
        erniebot.access_token = aistudio_access_token
        self.llm_model = "ernie-3.5"
        self.historical_messages = []
        self.function_map = function_map
        self._init_system_message(system_prompt)

    def _init_system_message(self, system_message: str):
        self.historical_messages = [
            {"role": "user", "content": system_message},
            {"role": "assistant", "content": "好的，我的男盆友",}
        ]

    def chat(self, user_message: str, on_func_call=None, on_func_complete=None, on_bot_reply=None):
        # 记录用户消息
        self.historical_messages.append({"role": "user", "content": user_message})
        # print(self.historical_messages)
        response = erniebot.ChatCompletion.create(
            model=self.llm_model, 
            messages=self.historical_messages,
            functions=[func['func_desc'] for func in self.function_map],
            stream=False
        )
        
        if hasattr(response, 'function_call'):

            function_call = response['function_call']
            # 记录bot要求函数调用消息
            self.historical_messages.append({
                "role": "assistant",
                "content": None,
                "function_call": function_call
                })
            func_name = function_call['name']
            func_args = json.loads(function_call['arguments'])
            bot_thoughts = function_call['thoughts']
            print(f"模型思考过程：{bot_thoughts}")

            if on_func_call:
                on_func_call(func_name, func_args, bot_thoughts) # 回调函数
                

            func_result = self.call_function(func_name, func_args)
            if on_func_complete:
                on_func_complete(func_name, func_args, func_result) # 回调函数

            # 记录函数调用结果消息
            self.historical_messages.append({"role": "function", "name": func_name, "content": func_result})
            response = erniebot.ChatCompletion.create(
                model=self.llm_model, 
                messages=self.historical_messages,
                functions=[func['func_desc'] for func in self.function_map],
                stream=False
            )
            bot_reply = response.get_result()
            if "|" not in bot_reply:
                bot_reply = "CALM|" + bot_reply
            assert not hasattr(response, 'function_call')
            self.historical_messages.pop() # 删除函数调用结果消息
            self.historical_messages.pop() # 删除函数调用消息
            self.historical_messages.append({
                "role": "assistant",
                "content": bot_reply
            })
            # return bot_reply

        else:
            bot_reply = response.get_result()
            self.historical_messages.append({"role": "assistant", "content": bot_reply})
        
        last_message = self.historical_messages[-1]
        if on_bot_reply:
            on_bot_reply(last_message["content"])

        # return bot_reply
    
    def call_function(self, func_name, func_args):
        for func in self.function_map:
            if func['func_desc']['name'] == func_name:
                func_result = func['func_pointer'](**func_args)
                func_result = json.dumps(func_result, ensure_ascii=False)
                assert func_result is not None
                # break
                return func_result
        else:
            return "未找到对应函数，调用失败"

        

# func_get_weather = {
#     "name": "get_weather",
#     "description": "get weather",
#     "parameters": {
#         "type": "object",
#         "required": ["city"],
#         "properties": {"city": {"type": "string"}},
#     },
# }

# def funcimp_get_weather(city):
#     return {
#         "city": city,
#         "weather": "晴",
#         "temperature": "25°C",
#         "wind": "东风3级",
#         "humidity": "50%",
#     }



# FUNC_MAP = [
#     {
#         "func_desc": func_get_weather,
#         "func_pointer": funcimp_get_weather,
#     },
# ]

if __name__ == "__main__":
    from FuncDefine import BOT_SYSTEM_PROMPT, FUNC_MAP
    aistudio_access_token = "8c1f786a2a69d12d76088175a0d401c160a41cf0"
    bot = ERNIE_API(aistudio_access_token, BOT_SYSTEM_PROMPT, FUNC_MAP)
    # bot_reply = bot.chat("今天北京的天气")
    # print(bot_reply)
    # bot_reply = bot.chat("今天天津的天气")
    # print(bot_reply)
    bot_reply = bot.chat("你好呀")
    print(bot_reply)
    # bot_reply = bot.chat("我们分手吧")
    # print(bot_reply)
    # bot_reply = bot.chat("记住明天早晨8点去海滩玩")
    # print(bot_reply)
    bot_reply = bot.chat("我们去海滩吧")
    print(bot_reply)

    

